import requests
import json
import time
import streamlit as st

# API 配置
BOT_ID = "7450905437840556073"
API_TOKEN = "pat_Hs2EnvB5JmAShsIMlzVHH9SjAnIcO71KBiYaoevIrP7FDZJ9uMzxcKHYxg6tG8T4"
BASE_URL = "https://api.coze.cn/v3"
WORKFLOW_API_TOKEN = "pat_Hs2EnvB5JmAShsIMlzVHH9SjAnIcO71KBiYaoevIrP7FDZJ9uMzxcKHYxg6tG8T4"
WORKFLOWS = {
    "text": "7448626255131705344",  # 文本处理工作流 ID
    "image": "7448825195676680201",  # 图片处理工作流 ID
    "image_creation": "7448829527528144905",  # 图片创作工作流 ID
    "video": "7448834266979614758",  # 视频处理工作流 ID
    "url": "7448836983155114020",  # URL 处理工作流 ID
}


# 发送聊天请求
def send_chat_request(question):
    url = f"{BASE_URL}/chat"
    headers = {
        "Authorization": f"Bearer {API_TOKEN}",
        "Content-Type": "application/json"
    }
    data = {
        "bot_id": BOT_ID,
        "user_id": "123456789",  # 可以使用任意用户ID
        "stream": False,
        "auto_save_history": True,
        "additional_messages": [
            {
                "role": "user",
                "content": question,
                "content_type": "text"
            }
        ]
    }

    response = requests.post(url, headers=headers, json=data)
    return response.json()


# 获取聊天消息
def get_chat_messages(chat_id, conversation_id):
    url = f"{BASE_URL}/chat/message/list?chat_id={chat_id}&conversation_id={conversation_id}"
    headers = {
        "Authorization": f"Bearer {API_TOKEN}",
        "Content-Type": "application/json"
    }

    response = requests.get(url, headers=headers)
    return response.json()


# 调用 Coze 工作流 API 的辅助函数
def call_workflow(workflow_id, user_description, param_key="input"):
    headers = {
        "Authorization": f"Bearer {WORKFLOW_API_TOKEN}",
        "Content-Type": "application/json"
    }
    data = {
        "workflow_id": workflow_id,
        "parameters": {
            "user_id": "12345",
            "user_name": "George",
            param_key: user_description  # 根据不同的工作流类型使用不同的参数键
        }
    }
    response = requests.post("https://api.coze.cn/v1/workflow/run", headers=headers, json=data)
    if response.status_code == 200:
        result = response.json()
        if result["code"] == 0:
            return result["data"]  # 返回工作流的执行结果
        else:
            return f"工作流执行失败：{result['msg']}"  # 返回错误信息
    elif response.status_code == 504:
        return "工作流执行超时。"
    else:
        return f"API 调用失败：状态码 {response.status_code}"


# 根据输入信息判断调用哪个工作流函数，简单示例，可根据实际情况完善逻辑
def determine_workflow(input_text):
    if "image" in input_text.lower():
        return "image"
    elif "video" in input_text.lower():
        return "video"
    elif "url" in input_text.lower():
        return "url"
    elif "image creation" in input_text.lower():
        return "image_creation"
    else:
        return "text"


def main():
    st.title("Coze API交互应用")
    input_text = st.text_input("请输入内容:")
    if st.button("提交"):
        with st.spinner("正在处理，请稍候..."):
            workflow_type = determine_workflow(input_text)
            workflow_id = WORKFLOWS[workflow_type]
            workflow_result = call_workflow(workflow_id, input_text)
            question = f"{input_text}\n参考内容：{workflow_result}"
            # 发送聊天请求
            chat_response = send_chat_request(question)
            if "code" in chat_response and chat_response["code"] == 0 and "data" in chat_response:
                chat_id = chat_response["data"].get("id")
                conversation_id = chat_response["data"].get("conversation_id")
                if chat_id and conversation_id:
                    for attempt in range(5):
                        messages = get_chat_messages(chat_id, conversation_id)
                        if "code" in messages and messages["code"] == 0 and "data" in messages and messages["data"]:
                            for message in messages["data"]:
                                if message["role"] == "assistant" and message["type"] == "answer":
                                    additional_info = f" 调用的工作流: {workflow_type}, 工作流id: {workflow_id}, 会话id: {conversation_id}, 智能体id: {BOT_ID}"
                                    st.success(f"智能体的回答：{message['content']}{additional_info}")
                                    return
                            st.warning("未找到智能体的回答，正在重试...")
                            time.sleep(5)
                        else:
                            st.warning("获取聊天消息失败或返回格式不正确，正在重试...")
                            time.sleep(5)
                    st.error("所有重试都失败，无法获取智能体的回答")
                else:
                    st.error("chat_id 或 conversation_id 未在响应中找到")
            else:
                st.error("发送聊天请求失败或返回格式不正确")


if __name__ == "__main__":
    main()